from keras import layers
from keras.layers import Input, Dense, Activation, ZeroPadding2D, BatchNormalization, Flatten, Conv2D,concatenate
from keras.layers import AveragePooling2D, MaxPooling2D, Dropout, GlobalMaxPooling2D, GlobalAveragePooling2D
from keras.models import Model

from keras.utils import plot_model

import numpy as np


def WorkerModel(worker_shape,slot_shape,backlog_shape):
    """
    Implementation of the HappyModel.

    Arguments:
    input_shape -- shape of the images of the dataset

    Returns:
    model -- a Model() instance in Keras
    """

    ### START CODE HERE ###
    # Feel free to use the suggested outline in the text above to get started, and run through the whole
    # exercise (including the later portions of this notebook) once. The come back also try out other
    # network architectures as well.
    # Define the input placeholder as a tensor with shape input_shape. Think of this as your input image!
    W_input = Input(worker_shape ,name="work_input")
    X1 = Flatten()(W_input)
    X1 = Activation('relu')(X1)
    X1 = Dense(30, activation='tanh', name='fc1')(X1)

    S_input = Input(slot_shape,name="slot_input")
    X2 = Flatten()(S_input)
    # X2 = Activation('relu')(X2)
    X2 = Dense(8, activation='tanh', name='fc2')(X2)

    B_input = Input(backlog_shape,name="backlog")
    X3 = Flatten()(B_input)
    # X3 = Activation('relu')(X3)
    X3 = Dense(4, activation='tanh', name='fc3')(X3)

    Y = concatenate([X1,X2,X3])
    # print(Y.shape)
    Y = Dense(6,activation='relu', name='fc')(Y)
    Y = Activation("softmax")(Y)
    model = Model(inputs=[W_input,S_input,B_input], outputs=Y, name="ScheduleModel")


    # X = ZeroPadding2D(padding=(3 ,3))(X_input)
    # CONV -> BN -> RELU Block applied to X
    # X = Conv2D(8 ,(2 ,2) ,strides=(1 ,1) ,name="conv0")(W_input)
    # # X = BatchNormalization(axis = 3, name = 'bn0')(X)
    #
    # X = Activation('relu')(X)
    #
    # # MAXPOOL
    # X = MaxPooling2D((2 ,2) ,name="max_pool")(X)
    #
    # # FLATTEN X (means convert it to a vector) + FULLYCONNECTED
    # X = Flatten()(X)
    # X = Dense(20, activation='sigmoid', name='fc')(X)
    #
    # model = Model(inputs=W_input ,outputs=X ,name="HappyModel")
    ### END CODE HERE ###

    return model



worker = np.array([[1., 1., 1., 1., 0., 0., 0., 0.],
                   [1., 1., 1., 1., 1., 1., 1., 0.],
                   [2., 2., 2., 2., 2., 0., 0., 0.],
                   [0., 0., 0., 0., 0., 0., 0., 0.],
                   [0., 0., 0., 0., 0., 0., 0., 0.],
                   [3., 3., 3., 3., 0., 0., 0., 0.],
                   [0., 0., 0., 0., 0., 0., 0., 0.],
                   [0., 0., 0., 0., 0., 0., 0., 0.],
                   [4., 4., 0., 0., 0., 0., 0., 0.],
                   [0., 0., 0., 0., 0., 0., 0., 0.]],dtype=np.int8)

slot = np.array([[3., 3., 3., 3., 3., 0., 0., 0.],
                 [3., 3., 3., 0., 0., 0., 0., 0.],
                 [1., 1., 1., 1., 1., 1., 1., 0.],
                 [3., 3., 3., 3., 3., 0., 0., 0.],
                 [0., 0., 0., 0., 0., 0., 0., 0.]],dtype=np.int8)

backlog = np.array([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]],dtype=np.int8)

wm = WorkerModel(worker.shape,slot.shape,backlog.shape)
wm.summary()

plot_model(wm, to_file='WorkerModel.png')